#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Quick checks for frame folders:
- Count JPGs per folder
- Verify %06d.jpg sequential naming and report missing/duplicates
- Print first/last image shape if OpenCV is available

Usage:
  python scripts/verify_frames.py --root Z:/Yolo训练集/03_frames_raw
  python scripts/verify_frames.py --folder Z:/Yolo训练集/03_frames_raw/20250909_locA_cam01
"""

import argparse
import os
import re
import glob


def list_folders(root: str):
    return [p for p in glob.glob(os.path.join(root, '*')) if os.path.isdir(p)]


def seq_stats(folder: str):
    files = sorted(glob.glob(os.path.join(folder, '*.jpg')))
    pat = re.compile(r"(\d{6})\.jpg$")
    idx = []
    for f in files:
        m = pat.search(os.path.basename(f))
        if m:
            idx.append(int(m.group(1)))
    missing = []
    if idx:
        start, end = min(idx), max(idx)
        s = set(idx)
        missing = [i for i in range(start, end + 1) if i not in s]
    return len(files), missing


def image_shape(path: str):
    try:
        import cv2  # type: ignore
        im = cv2.imread(path)
        if im is None:
            return None
        h, w = im.shape[:2]
        return (h, w)
    except Exception:
        return None


def check_folder(folder: str):
    n, missing = seq_stats(folder)
    first = os.path.join(folder, '000001.jpg')
    last = os.path.join(folder, f"{n:06d}.jpg") if n > 0 else None
    shp1 = image_shape(first) if os.path.exists(first) else None
    shp2 = image_shape(last) if last and os.path.exists(last) else None
    print(f"[Folder] {folder}")
    print(f"  count: {n}")
    if shp1:
        print(f"  first: 000001.jpg shape={shp1}")
    if shp2:
        print(f"  last : {os.path.basename(last)} shape={shp2}")
    if missing:
        print(f"  missing indices: {len(missing)} (showing first 20): {missing[:20]}")
    else:
        print("  sequence: OK")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--root', type=str, help='Root dir that contains <video_id>/ folders')
    ap.add_argument('--folder', type=str, help='Check a single folder')
    args = ap.parse_args()

    if args.folder:
        check_folder(args.folder)
        return
    if not args.root:
        ap.error('--root or --folder is required')
    for d in list_folders(args.root):
        check_folder(d)


if __name__ == '__main__':
    main()

